import numpy as np
from OrderDispatch.OrderDispatchRL.Planning.KM_Algorithm import KM_Algorithm
import random
from copy import deepcopy
import os
from OrderDispatch.OrderDispatchRL.TimeEstimate.TimeEstimaterV1 import TimeEstimaterV1
from General.MapLoader import MapLoader
import pandas as pd

class Planner:
    """
        该强化学习订单分配算法分为两个步骤，learning和planning
        2）planning：
        规划阶段，目的是完成空闲机器人与待分配订单之间的匹配，根据learning阶段得到的状态价值判断机器人完成某订单能获得的收益
        规划阶段将订单分配问题建模为空闲机器人与待分配订单之间的二分图最大匹配问题，边权重为对应的优势函数值（通过奖励与状态价值计算得到），再使用KM算法求解该二分图匹配问题
    """
    def __init__(self):
        self.Vs = np.load(os.path.join(os.path.dirname(__file__), '../data/Vs.npy'))
        self.totalTime = self.Vs.shape[0]
        self.Region = self.Vs.shape[1]

        self.result = []
        self.OrdNum = 10  # The number of Orders each period
        self.DriNum = 20  # The number of dirvers,it should be not greater than OrdNum

        self.Orders = []
        # 待分配订单列表，例：self.Orders = [[0, (0, 4), 0, 70, (1, 1), 0], [1, (0, 4), 0, 70, (1, 21), 0]]
        # [订单编号, (开始时间, 开始位置), 是否已被分配, 订单价值, (结束时间, 结束位置), 未匹配时间]
        self.Drivers = []
        # 空闲机器人列表，例：self.Drivers = [[0, (0, 11), 0, 0]]
        # [机器人编号, (当前时间, 当前所在位置), 是否已被分配, 当前收入]
        self.finalresult = np.full((self.DriNum, self.totalTime), -1, dtype=int)
        self.OrdNo = -1
        self.flag = 0  # Default model(DirNum < OrdNum)

        self.weight_params = {'weight_1': 0.4, 'weight_2': 0.3, 'weight_3': 0.3}
        self.code_index_map = pd.read_excel(
            os.path.join(os.path.dirname(__file__), '../data/node_index_with_cluster.xlsx'))

    def transfer_timestamp(self, int_time):
        return int_time % self.totalTime

    def get_region_by_node_index(self, node_index):
        return self.code_index_map.loc[self.code_index_map['index'] == node_index, 'cluster'].iloc[0]

    def calculate_order_value(self, time1, time2, time3):
        """计算订单价值/动作奖励"""
        time_diff_1 = time1 / 60.0
        time_diff_2 = time2 / 60.0
        time_diff_3 = time3 / 60.0

        # 使用传入的权重参数计算加权平均值
        return (self.weight_params['weight_1'] * time_diff_1) + \
               (self.weight_params['weight_2'] * time_diff_2) + \
               (self.weight_params['weight_3'] * time_diff_3)


    def Calc_AdvFuc(self, order_list, robot_list):
        """计算优势函数值（二分图中的边权重）"""
        time_estimater = TimeEstimaterV1()
        map_loader = MapLoader()
        map_loader.create_graph_from_files()

        global Af
        if (len(order_list) > len(robot_list)):
            self.flag = 0
            Af = np.zeros((len(order_list), len(order_list)))  # Advantage function
        if (len(robot_list) > len(order_list)):
            self.flag = 1  # Special model
            Af = np.zeros((len(robot_list), len(robot_list)))  # Advantage function
        for i in range(0, len(robot_list)):
            for j in range(0, len(order_list)):
                # t = abs(robot_list[i][1][1] - order_list[j][1][1])  # 测试
                # t = freeOrd[j][4][0] - freeOrd[j][1][0]  #Time gap
                print(i, j)
                # position1为机器人当前位置
                position1 = map_loader.get_node_code_by_index(robot_list[i].current_position)
                # position2为订单取货点位置
                position2 = map_loader.get_node_code_by_index(order_list[j].pickup_point)
                # position3为订单送货点位置
                position3 = map_loader.get_node_code_by_index(order_list[j].delivery_point)
                # time1为订单等待时间
                time1 = (order_list[j].start_time - order_list[j].create_time)
                # time2为预测的机器人到达取货点时间
                time2 = time_estimater.predict(pd.DataFrame({
                        'create_hour': [order_list[j].start_time / 60 / 24],
                        'create_weekday': [1],
                        'start_node': [position1],
                        'end_node': [position2]
                    }))[0]
                # time3为预测的机器人到达送货点时间
                time3 = time_estimater.predict(pd.DataFrame({
                    'create_hour': [order_list[j].start_time / 60 / 24],
                    'create_weekday': [1],
                    'start_node': [position2],
                    'end_node': [position3]
                }))[0]
                start_timestamp = self.transfer_timestamp(order_list[j].start_time)
                start_region = self.get_region_by_node_index(robot_list[i].current_position)
                end_timestamp = self.transfer_timestamp(order_list[j].start_time + int(time2 + time3))
                end_region = self.get_region_by_node_index(order_list[j].delivery_point)
                value = self.calculate_order_value(time1, time2, time3)
                # 计算优势函数值，Af[i][j]表示机器人i执行订单j的优势函数值，值越大表示机器人执行该订单收益越大
                Af[i][j] = round(
                    pow(0.9, end_timestamp - start_timestamp) * self.Vs[end_timestamp][end_region] - self.Vs[start_timestamp][start_region] + value, 2)
                if (Af[i][j] < 0):  # The result is bad,make it zero!
                    Af[i][j] = 0
            # print('('+str(i)+str(j)+'):'+str(V_s[freeOrd[j][4][0]][freeOrd[j][4][1]])+'-'+str(V_s[freeOrd[i][1][0]][freeOrd[i][1][1]])+'+'+str(freeOrd[j][3])+'='+str(Af[i][j]))
        return Af

    def dispatch_order(self, order_list, robot_list):
        """实现订单分配，计算af优势函数矩阵，将其作为二分图的边权重，使用KM算法求解"""
        af = self.Calc_AdvFuc(order_list, robot_list)
        KM = KM_Algorithm(af)
        KM.Kuh_Munkras()
        # print('分配结果')
        # print(KM.getResult())
        return self.get_dispatch_results(KM.getResult(), order_list, robot_list)

    def get_dispatch_results(self, km_result, order_list, robot_list):
        """得到分配结果"""
        re = km_result
        matching_pairs = []
        if self.flag == 0:  # default model ,dir < ord
            for i in range(len(order_list)):
                if re[i] < len(robot_list):
                    robot = robot_list[re[i]]
                    order = order_list[i]

                    # 保存匹配对 (driver, order)
                    matching_pairs.append((robot, order))

                    # print(f'({driver_idx}, {order_idx})')

        if self.flag == 1:  # special model, dri > ord
            for i in range(len(order_list)):
                robot = robot_list[re[i]]
                order = order_list[i]
                # 保存匹配对 (driver, order)
                matching_pairs.append((robot, order))

                # print(f'({driver_idx}, {order_idx})')
        return matching_pairs


    def OrdersProductor(self, timestap):
        # 测试
        for i in range(self.OrdNum):  # Producing 20 orders
            self.OrdNo += 1
            t1 = timestap  # start time should be same with timestap
            g1 = random.randint(0, self.Region - 1)
            t2 = t1 + random.randint(1, self.totalTime - t1 - 1)
            g2 = random.randint(0, self.Region - 1)
            self.Orders.append([self.OrdNo, (t1, g1), 0, random.randint(10, 100), (t2, g2), 0])
        return self.Orders

    def DriverProductor(self):
        # 测试
        for i in range(self.DriNum):
            g = random.randint(0, self.Region - 1)  # Choose a place randomly
            self.Drivers.append([i, (0, g), 0, 0])  # No,state,used,money

        return self.Drivers

    def getFreeOrd(self, timestap):
        bk = deepcopy(self.Orders)
        for sub in bk:
            if (sub[5] > 1 or sub[2] == 1 or sub[1][0] < timestap):
                # print("delete invaild and mathed Order No.: "+str(sub[0]))
                self.Orders.remove(sub)  # Remove the order that is unmatched for long time(3 timestap)
                continue
        return self.Orders

    def getFreeDri(self, timestap):
        subDri = []
        for sub in self.Drivers:
            if (sub[2] == 0):
                subDri.append(sub)

        return subDri


if __name__ == '__main__':
    # test
    PL = Planner()
    PL.DriverProductor()  # Create drivers
    time = 0
    while (time < 10):
        # RL.OrdersProductor()

        PL.OrdersProductor(time)  # timestap

        fo = PL.getFreeOrd(time)
        fd = PL.getFreeDri(time)
        print("time " + str(time) + ' has drivers:')
        print(fd)
        print("time " + str(time) + ' has orders:')
        print(fo)

        PL.dispatch_order(fo, fd)
        time = time + 1